#include "rsa.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/**
 * @brief calculate a power with a module.
 *
 * We need to do this in order to keep usable bitlength number.
 * If we do not, we will overflow.
 *
 * @param a Number to compute a power
 * @param b Power to compute on a
 * @param m Modulus to apply to the calcul
 *
 * @return 0 in case of error, the result of (a ^ b) % m otherwise
 */
static t_uint64       modular_exponentiation(t_uint64 a, t_uint64 b, t_uint64 m)
{
    t_uint64          p;

    p = 1;
    while (b > 0)
        {
            if ((b & 1) != 0)
                p = (p * a) % m;
            a = (a * a) % m;
            b >>= 1;
        }
    return (p);
}

/**
 * @brief Encrypt a string
 *
 * In order to encrypt a string, we encrypt char by char. We will see later
 * how to have a better way to encrypt a string.
 *
 * @see encrypt_char
 *
 * @note The encrypted message is allocated, you need to pass it to the dercrypt
 *       function or to free it yourself.
 *
 * @param pub Public Key
 * @param msg Message string to encrypt
 *
 * @return Encrypted string (array)
 */
t_cryptedMessage        encrypt_message(t_publicKey *pub, char *msg)
{
    t_cryptedMessage    c;
    t_uint64            i;

    c = malloc(sizeof(t_cryptedMessage) * (strlen(msg) + 1));
    for (i = 0; i < strlen(msg); i++)
        c[i] = modular_exponentiation(msg[i], pub->e, pub->n);
    c[i] = 0;
    return (c);
}

/**
 * @brief Decrypt a string
 *
 * In order to decrypt a string, we decrypt char by char. We will see later
 * how to have a better way to decrypt a string.
 *
 * @see decrypt_char
 *
 * @note The decrypted string is allocated, you need to free it after.
 * @note The encrypted message is going to be freed.
 *
 * @param priv Private Key
 * @param c Encrypted string (array)
 *
 * @return Decrypted string
 *
 */
char            *decrypt_message(t_privateKey *priv, t_cryptedMessage c)
{
    char        *str;
    t_uint64    i;

    for (i = 0; c[i] != 0; i++)
        ;
    str = malloc(sizeof(char) * (i + 1));
    for (i = 0; c[i] != 0; i++)
        str[i] = modular_exponentiation(c[i], priv->d, priv->n);
    str[i] = 0;
    free(c);
    return (str);
}

/**
 * @brief Encrypt a character.
 *
 * In order to encrypt a char c, we take the ascii value
 * and we apply the formula : c ^ e % n
 *
 * @param pub Public Key
 * @param c char to encrypt
 *
 * @return Encrypted char
 */
t_cryptedChar   encrypt_char(t_publicKey *pub, char c)
{
    return (modular_exponentiation(c, pub->e, pub->n));
}

/**
 * @brief Decrypt a character.
 *
 * In order to decrypt a char c, we use the formula: c ^ d % n
 *
 * @param priv Private Key
 * @param c char to decrypt
 *
 * @return Decrtyped char
 */
char            decrypt_char(t_privateKey *priv, t_cryptedChar c)
{
    return (modular_exponentiation(c, priv->d, priv->n));
}

int                     main()
{
    t_publicKey         pub;
    t_privateKey        priv;
    t_cryptedMessage    cMsg;
    t_cryptedChar       c;
    char                *str;

    if (generate_keys(&pub, &priv, 0, 0) == -1)
        {
            fprintf(stderr, "An error occured while generating the keys\n");
            return (EXIT_FAILURE);
        }
    printf("p=%lu, q=%lu, n=%lu, phi(n) = %lu, e=%lu, d=%lu\n", priv.p, priv.q, pub.n, priv.phi_n, pub.e, priv.d);

    printf("c = m^e %% n, c = %d^%lu %% %lu\n", 'o', pub.e, pub.n);

    c = encrypt_char(&pub, 'o');
    printf("c: %lu\n", c);
    printf("m = c^d %% n, c = %lu^%lu %% %lu\n", c, priv.d, pub.n);
    printf("m: %c\n", decrypt_char(&priv, c));

    cMsg = encrypt_message(&pub, "hello world!");
    str = decrypt_message(&priv, cMsg);
    printf("%s\n", str);
    free(str);
    return (0);
}
